import logging

import dash
import os
import requests
import flask.cli
from IPython.core.display import HTML
from retrying import retry
import io
import re
import sys
import inspect
import traceback
import threading
import warnings
import queue

from IPython import get_ipython
from IPython.display import IFrame, display
from IPython.core.ultratb import FormattedTB
from ansi2html import Ansi2HTMLConverter
import uuid

from werkzeug.serving import make_server

from .comms import _dash_comm, _jupyter_config, _request_jupyter_config


def _get_skip(error: Exception):
    tb = traceback.format_exception(type(error), error, error.__traceback__)
    skip = 0
    for i, line in enumerate(tb):
        if "%% callback invoked %%" in line:
            skip = i + 1
            break
    return skip


class JupyterDash(dash.Dash):
    """A Dash subclass for developing Dash apps interactively in Jupyter.

    :param server_url:  The base URL that the app will be served at, from the
        perspective of the client. If not specified, will default to the host argument
        passed to the ``run_server`` method.

    See parent docstring for additional parameters
    """
    default_mode = 'external'
    default_requests_pathname_prefix = os.environ.get(
        "DASH_REQUESTS_PATHNAME_PREFIX", None
    )
    default_server_url = None
    _in_ipython = get_ipython() is not None
    _in_colab = "google.colab" in sys.modules
    _token = str(uuid.uuid4())

    _servers = {}

    @classmethod
    def infer_jupyter_proxy_config(cls):
        """
        Infer the current Jupyter server configuration. This will detect
        the proper request_pathname_prefix and server_url values to use when
        displaying Dash apps.  When the jupyter_server_proxy Python package is
        installed, all Dash requests will be routed through the proxy.

        Requirements:

        In the classic notebook, this method requires the `jupyter_dash` nbextension
        which should be installed automatically with the installation of the
        jupyter-dash Python package. You can see what notebook extensions are installed
        by running the following command:
            $ jupyter nbextension list

        In JupyterLab, this method requires the `jupyterlab-dash` labextension. This
        extension should be installed automatically with the installation of the
        jupyter-dash Python package, but JupyterLab must be allowed to rebuild before
        the extension is activated (JupyterLab should automatically detect the
        extension and produce a popup dialog asking for permission to rebuild). You can
        see what JupyterLab extensions are installed by running the following command:
            $ jupyter labextension list
        """
        if not JupyterDash._in_ipython or JupyterDash._in_colab:
            # No op when not running in a Jupyter context or when in Colab
            return
        else:
            # Assume classic notebook or JupyterLab
            _request_jupyter_config()

    def __init__(self, name=None, server_url=None, **kwargs):
        """"""
        # Strip unsupported properties and warn
        if JupyterDash._in_colab:
            unsupported_colab_props = [
                'requests_pathname_prefix',
                'routes_pathname_prefix',
                'url_base_pathname'
            ]
            for prop in unsupported_colab_props:
                if prop in kwargs:
                    kwargs.pop(prop)
                    warnings.warn(
                        "The {prop} argument is ignored when running in Colab".format(
                            prop=prop
                        )
                    )

        # Call superclass constructor
        super(JupyterDash, self).__init__(name=name, **kwargs)

        if not JupyterDash._in_ipython:
            # Nothing else to do when not running in a Jupyter context
            return

        # See if jupyter_server_proxy is installed
        try:
            import jupyter_server_proxy
            self._server_proxy = True
        except Exception:
            self._server_proxy = False

        self._traceback = None

        if ('base_subpath' in _jupyter_config and self._server_proxy and
                JupyterDash.default_requests_pathname_prefix is None):
            JupyterDash.default_requests_pathname_prefix = (
                _jupyter_config['base_subpath'].rstrip('/') + '/proxy/{port}/'
            )

        if ('server_url' in _jupyter_config and self._server_proxy and
                JupyterDash.default_server_url is None):
            JupyterDash.default_server_url = _jupyter_config['server_url']

        self._input_pathname_prefix = kwargs.get('requests_pathname_prefix', None)

        # Infer server_url
        if server_url is None:
            domain_base = os.environ.get('DASH_DOMAIN_BASE', None)
            if domain_base:
                # Dash Enterprise sets DASH_DOMAIN_BASE environment variable
                server_url = 'https://' + domain_base
        elif JupyterDash._in_colab:
            warnings.warn("The server_url argument is ignored when running in Colab")
            server_url = None

        self.server_url = server_url

        # Register route that we can use to poll to see when server is running
        @self.server.route('/_alive_' + JupyterDash._token, methods=['GET'])
        def alive():
            return 'Alive'

        self.server.logger.disabled = True
        self._exception_handling_added = False

    def run(
            self,
            mode=None, width="100%", height=650, inline_exceptions=None,
            **kwargs
    ):
        """
        Serve the app using flask in a background thread. You should not run this on a
        production server, use gunicorn/waitress instead.

        :param mode: Display mode. One of:
            ``"external"``: The URL of the app will be displayed in the notebook
                output cell. Clicking this URL will open the app in the default
                web browser.
            ``"inline"``: The app will be displayed inline in the notebook output cell
                in an iframe.
            ``"jupyterlab"``: The app will be displayed in a dedicate tab in the
                JupyterLab interface. Requires JupyterLab and the `jupyterlab-dash`
                extension.
        :param width: Width of app when displayed using mode="inline"
        :param height: Height of app when displayed using mode="inline"
        :param inline_exceptions: If True, callback exceptions are displayed inline
            in the the notebook output cell. Defaults to True if mode=="inline",
            False otherwise.
        :param kwargs: Additional keyword arguments to pass to the superclass
            ``Dash.run_server`` method.
        """
        # Get superclass run_server method
        if hasattr(dash.Dash, "run"):
            super_run_server = super(JupyterDash, self).run
        else:
            super_run_server = super(JupyterDash, self).run_server

        if not JupyterDash._in_ipython:
            # If not in IPython context, call run run_server synchronously
            super_run_server(**kwargs)
            return

        # Get host and port
        host = kwargs.pop("host", os.getenv("HOST", "127.0.0.1"))
        port = int(kwargs.pop("port", os.getenv("PORT", "8050")))

        # Validate / infer display mode
        if JupyterDash._in_colab:
            valid_display_values = ["inline", "external"]
        else:
            valid_display_values = ["jupyterlab", "inline", "external"]

        if mode is None:
            mode = JupyterDash.default_mode
        elif not isinstance(mode, str):
            raise ValueError(
                "The mode argument must be a string\n"
                "    Received value of type {typ}: {val}".format(
                    typ=type(mode), val=repr(mode)
                )
            )
        else:
            mode = mode.lower()
            if mode not in valid_display_values:
                raise ValueError(
                    "Invalid display argument {display}\n"
                    "    Valid arguments: {valid_display_values}".format(
                        display=repr(mode), valid_display_values=valid_display_values
                    )
                )

        # Infer inline_exceptions and ui
        if inline_exceptions is None:
            inline_exceptions = mode == "inline"

        # Terminate any existing server using this port
        old_server = self._servers.get((host, port))
        if old_server:
            old_server.shutdown()
            del self._servers[(host, port)]

        # Configure pathname prefix
        requests_pathname_prefix = self.config.get('requests_pathname_prefix', None)
        if self._input_pathname_prefix is None:
            requests_pathname_prefix = self.default_requests_pathname_prefix

        if requests_pathname_prefix is not None:
            requests_pathname_prefix = requests_pathname_prefix.format(port=port)
        else:
            requests_pathname_prefix = '/'
        # low-level setter to circumvent Dash's config locking
        # normally it's unsafe to alter requests_pathname_prefix this late, but
        # Jupyter needs some unusual behavior.
        dict.__setitem__(self.config, "requests_pathname_prefix", requests_pathname_prefix)

        # Compute server_url url
        if self.server_url is None:
            if JupyterDash.default_server_url:
                server_url = JupyterDash.default_server_url.rstrip('/')
            else:
                server_url = 'http://{host}:{port}'.format(host=host, port=port)
        else:
            server_url = self.server_url.rstrip('/')

        dashboard_url = "{server_url}{requests_pathname_prefix}".format(
            server_url=server_url, requests_pathname_prefix=requests_pathname_prefix
        )

        # Default the global "debug" flag to True
        debug = kwargs.pop('debug', True)

        # Enable supported dev tools
        if debug:
            for k in [
                'dev_tools_silence_routes_logging',
                'dev_tools_props_check',
                'dev_tools_serve_dev_bundles',
                'dev_tools_prune_errors'
            ]:
                if k not in kwargs:
                    kwargs[k] = True

            # Enable dev tools by default unless app is displayed inline
            if 'dev_tools_ui' not in kwargs:
                kwargs['dev_tools_ui'] = mode != "inline"

            if 'dev_tools_hot_reload' not in kwargs:
                # Enable hot-reload by default in "external" mode. Enabling in inline or
                # in JupyterLab extension seems to cause Jupyter problems sometimes when
                # there is no active kernel.
                kwargs['dev_tools_hot_reload'] = mode == "external"

        # Set up custom callback exception handling
        self._config_callback_exception_handling(
            dev_tools_prune_errors=kwargs.get('dev_tools_prune_errors', True),
            inline_exceptions=inline_exceptions,
        )

        dev_tools_args = dict(
            debug=debug,
            dev_tools_ui=kwargs.pop("dev_tools_ui", None),
            dev_tools_props_check=kwargs.pop("dev_tools_props_check", None),
            dev_tools_serve_dev_bundles=kwargs.pop("dev_tools_serve_dev_bundles", None),
            dev_tools_hot_reload=kwargs.pop("dev_tools_hot_reload", None),
            dev_tools_hot_reload_interval=kwargs.pop("dev_tools_hot_reload_interval", None),
            dev_tools_hot_reload_watch_interval=kwargs.pop("dev_tools_hot_reload_watch_interval", None),
            dev_tools_hot_reload_max_retry=kwargs.pop("dev_tools_hot_reload_max_retry", None),
            dev_tools_silence_routes_logging=kwargs.pop("dev_tools_silence_routes_logging", None),
            dev_tools_prune_errors=kwargs.pop("dev_tools_prune_errors", None),
        )

        if len(kwargs):
            raise Exception(f"Invalid keyword argument: {list(kwargs.keys())}")

        self.enable_dev_tools(**dev_tools_args)

        # suppress warning banner printed to standard out
        flask.cli.show_server_banner = lambda *args, **kw: None

        # prevent partial import of orjson when it's installed and mode=jupyterlab
        # TODO: why do we need this? Why only in this mode? Importing here in
        # all modes anyway, in case there's a way it can pop up in another mode
        try:
            import orjson
        except ImportError:
            pass

        err_q = queue.Queue()

        server = make_server(
            host, port, self.server,
            threaded=True,
            processes=0
        )
        logging.getLogger("werkzeug").setLevel(logging.ERROR)

        @retry(
            stop_max_attempt_number=15,
            wait_exponential_multiplier=100,
            wait_exponential_max=1000
        )
        def run():
            try:
                server.serve_forever()
            except SystemExit:
                pass
            except Exception as error:
                err_q.put(error)
                raise error

        thread = threading.Thread(target=run)
        thread.daemon = True
        thread.start()

        self._servers[(host, port)] = server

        # Wait for server to start up
        alive_url = "http://{host}:{port}/_alive_{token}".format(
            host=host, port=port, token=JupyterDash._token
        )

        def _get_error():
            try:
                err = err_q.get_nowait()
                if err:
                    raise err
            except queue.Empty:
                pass

        # Wait for app to respond to _alive endpoint
        @retry(
            stop_max_attempt_number=15,
            wait_exponential_multiplier=10,
            wait_exponential_max=1000
        )
        def wait_for_app():
            _get_error()
            try:
                req = requests.get(alive_url)
                res = req.content.decode()
                if req.status_code != 200:
                    raise Exception(res)

                if res != "Alive":
                    url = "http://{host}:{port}".format(
                        host=host, port=port, token=JupyterDash._token
                    )
                    raise OSError(
                        "Address '{url}' already in use.\n"
                        "    Try passing a different port to run_server.".format(
                            url=url
                        )
                    )
            except requests.ConnectionError as err:
                _get_error()
                raise err

        try:
            wait_for_app()

            if JupyterDash._in_colab:
                self._display_in_colab(dashboard_url, port, mode, width, height)
            else:
                self._display_in_jupyter(dashboard_url, port, mode, width, height)
        except Exception as final_error:
            msg = str(final_error)
            if msg.startswith('<!'):
                display(HTML(msg))
            else:
                raise final_error

    def _display_in_colab(self, dashboard_url, port, mode, width, height):
        from google.colab import output
        if mode == 'inline':
            output.serve_kernel_port_as_iframe(port, width=width, height=height)
        elif mode == 'external':
            # Display a hyperlink that can be clicked to open Dashboard
            print("Dash app running on:")
            output.serve_kernel_port_as_window(port, anchor_text=dashboard_url)

    def _display_in_jupyter(self, dashboard_url, port, mode, width, height):
        if mode == 'inline':
            display(IFrame(dashboard_url, width, height))
        elif mode == 'external':
            # Display a hyperlink that can be clicked to open Dashboard
            print("Dash app running on {dashboard_url}".format(
                dashboard_url=dashboard_url
            ))
        elif mode == 'jupyterlab':
            # Update front-end extension
            _dash_comm.send({
                'type': 'show',
                'port': port,
                'url': dashboard_url,
            })

    def _config_callback_exception_handling(
            self, dev_tools_prune_errors, inline_exceptions
    ):

        @self.server.errorhandler(Exception)
        def _wrap_errors(error):
            """Install traceback handling for callbacks"""
            self._traceback = sys.exc_info()[2]

            # Compute number of stack frames to skip to get down to callback
            skip = _get_skip(error) if dev_tools_prune_errors else 0

            # Customized formatargvalues function so we can place function parameters
            # on separate lines
            original_formatargvalues = inspect.formatargvalues
            inspect.formatargvalues = _custom_formatargvalues
            try:
                # Use IPython traceback formatting to build colored ANSI traceback
                # string
                ostream = io.StringIO()
                ipytb = FormattedTB(
                    tb_offset=skip,
                    mode="Verbose",
                    color_scheme="Linux",
                    include_vars=True,
                    ostream=ostream
                )
                ipytb()
            finally:
                # Restore formatargvalues
                inspect.formatargvalues = original_formatargvalues

            # Print colored ANSI representation if requested
            ansi_stacktrace = ostream.getvalue()

            if inline_exceptions:
                print(ansi_stacktrace)

            # Use ansi2html to convert the colored ANSI string to HTML
            conv = Ansi2HTMLConverter(scheme="ansi2html", dark_bg=False)
            html_str = conv.convert(ansi_stacktrace)

            # Set width to fit 75-character wide stack trace and font to a size the
            # won't require a horizontal scroll bar
            html_str = html_str.replace(
                '<html>',
                '<html style="width: 75ch; font-size: 0.86em">'
            )

            # Remove explicit background color so Dash dev-tools can set background
            # color
            html_str = re.sub("background-color:[^;]+;", "", html_str)

            return html_str, 500

    def run_server(
            self,
            mode=None, width="100%", height=650, inline_exceptions=None,
            **kwargs
    ):
        self.run(
            mode=mode, width=width, height=height, inline_exceptions=inline_exceptions,
            **kwargs
        )


def _custom_formatargvalues(
        args, varargs, varkw, locals,
        formatarg=str,
        formatvarargs=lambda name: '*' + name,
        formatvarkw=lambda name: '**' + name,
        formatvalue=lambda value: '=' + repr(value)):

    """Copied from inspect.formatargvalues, modified to place function
    arguments on separate lines"""
    def convert(name, locals=locals,
                formatarg=formatarg, formatvalue=formatvalue):
        return formatarg(name) + formatvalue(locals[name])
    specs = []
    for i in range(len(args)):
        specs.append(convert(args[i]))
    if varargs:
        specs.append(formatvarargs(varargs) + formatvalue(locals[varargs]))
    if varkw:
        specs.append(formatvarkw(varkw) + formatvalue(locals[varkw]))

    result = '(' + ', '.join(specs) + ')'

    if len(result) < 40:
        return result
    else:
        # Put each arg on a separate line
        return '(\n    ' + ',\n    '.join(specs) + '\n)'
